#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A股股票智能筛选系统 - Baostock版本
使用Baostock免费API（稳定、无需token）
"""

import sys
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

import baostock as bs
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import List, Dict
import warnings
warnings.filterwarnings('ignore')


class ScreenerConfig:
    """筛选配置 - 进一步放宽"""
    # 基础条件
    MIN_PRICE = 5  # 降至7元
    MAX_PRICE = 60  # 提高至60元
    MIN_VOLUME = 10_000_000  # 1000万元（大幅降低）
    MIN_LIST_YEARS = 1
    MIN_MARKET_CAP = 10_00_000_000  # 10亿元（降低至10亿）
    
    # 技术面指标
    MIN_3M_RETURN = -20  # 大幅放宽
    MAX_3M_RETURN = 35  # 大幅放宽
    MAX_52W_DRAWDOWN = 50  # 放宽至50%
    MIN_TURNOVER = 0.2  # 降低至0.2%
    
    # 优先行业
    PRIORITY_INDUSTRIES = [
        '半导体', '集成电路', '芯片', '电子', '软件', '云计算', '人工智能', '计算机', '通信',
        '医药', '生物', 'CXO', '医疗', '器械',
        '锂电', '光伏', '储能', '新能源', '电池', '风电', '电力设备',
        '机械', '自动化', '军工', '航空', '航天',
        '消费', '食品', '饮料', '白酒', '家电', '纺织', '零售'
    ]
    
    OUTPUT_COUNT = 20


class BaostockAPI:
    """Baostock API封装"""
    
    def __init__(self):
        print("[*] 连接Baostock服务器...")
        lg = bs.login()
        if lg.error_code != '0':
            raise Exception(f"Baostock登录失败: {lg.error_msg}")
        print("[OK] 连接成功\n")
    
    def __del__(self):
        bs.logout()
    
    def get_stock_list(self) -> List[Dict]:
        """获取A股股票列表"""
        print("[*] 获取A股股票列表...")
        rs = bs.query_stock_basic(code_name="")  # 获取所有股票
        
        stocks = []
        while (rs.error_code == '0') & rs.next():
            row = rs.get_row_data()
            code = row[0]  # sh.600000 或 sz.000001
            name = row[1]  # 股票名称（修复：原来错用了row[2]即上市日期）
            
            # 过滤规则
            # 1. 跳过指数（代码为000001-000300等）
            clean_code = code.split('.')[1]
            if clean_code.startswith('000') and len(clean_code) == 6:
                continue
            
            # 2. 只保留沪深A股（6开头的沪市，0/3开头的深市）
            if not (clean_code.startswith('6') or clean_code.startswith('0') or clean_code.startswith('3')):
                continue
            
            # 3. 排除ST、退市股
            if 'ST' in name or '退' in name or '*' in name:
                continue
            
            # 4. 排除北交所
            if code.startswith('bj.'):
                continue
            
            stocks.append({
                '代码': clean_code,
                '名称': name,
                'baostock代码': code
            })
        
        print(f"[OK] 获取 {len(stocks)} 支A股\n")
        return stocks
    
    def get_stock_kline(self, baostock_code: str, days: int = 800) -> pd.DataFrame:
        """获取K线数据"""
        end_date = datetime.now().strftime('%Y-%m-%d')
        start_date = (datetime.now() - timedelta(days=days)).strftime('%Y-%m-%d')
        
        rs = bs.query_history_k_data_plus(
            baostock_code,
            "date,code,open,high,low,close,volume,amount,turn,pctChg",
            start_date=start_date,
            end_date=end_date,
            frequency="d",
            adjustflag="2"  # 前复权
        )
        
        data = []
        while (rs.error_code == '0') & rs.next():
            data.append(rs.get_row_data())
        
        if not data:
            return pd.DataFrame()
        
        df = pd.DataFrame(data, columns=rs.fields)
        # 转换数据类型
        for col in ['open', 'high', 'low', 'close', 'volume', 'amount', 'turn', 'pctChg']:
            df[col] = pd.to_numeric(df[col], errors='coerce')
        
        return df


class StockScreener:
    """股票筛选器"""
    
    def __init__(self):
        self.api = BaostockAPI()
        self.config = ScreenerConfig()
        self.fetch_time = datetime.now()
        self.stats = {
            '初始股票池': 0,
            '基础条件过滤': 0,
            '行业筛选': 0,
            '技术面筛选': 0,
            '最终入选': 0
        }
    
    def run(self) -> pd.DataFrame:
        """执行筛选"""
        print("=" * 80)
        print(f"A股智能筛选系统 - Baostock版本".center(80))
        print("=" * 80)
        print(f"数据获取时间: {self.fetch_time.strftime('%Y-%m-%d %H:%M:%S')}\n")
        
        # 1. 获取股票列表
        stock_list = self.api.get_stock_list()
        if not stock_list:
            print("[ERROR] 无法获取股票列表")
            return pd.DataFrame()
        
        self.stats['初始股票池'] = len(stock_list)
        
        # 2. 获取K线数据
        print("【阶段1】获取K线数据并筛选...")
        print(f"  股票列表: {len(stock_list)} 支")
        print(f"  处理范围: 前300支\n")
        
        all_data = []
        failed_count = 0
        process_count = min(len(stock_list), 300)
        
        for idx, stock in enumerate(stock_list[:process_count], 1):
            if idx % 50 == 0:
                print(f"  进度: {idx}/{process_count}, 成功: {len(all_data)}, 失败: {failed_count}")
            
            df_kline = self.api.get_stock_kline(stock['baostock代码'])
            if df_kline.empty:
                failed_count += 1
                continue
            if len(df_kline) < 60:
                failed_count += 1
                continue
            
            # 计算指标
            result = self._calculate_indicators(stock, df_kline)
            if result:
                all_data.append(result)
            else:
                failed_count += 1
        
        print(f"\n  处理完成: 成功 {len(all_data)} 支, 失败 {failed_count} 支")
        
        if not all_data:
            print("[ERROR] 未获取到有效数据")
            return pd.DataFrame()
        
        df = pd.DataFrame(all_data)
        print(f"[OK] 成功获取 {len(df)} 支股票数据\n")
        
        # 3. 基础筛选
        print("【阶段2】基础条件筛选...")
        print(f"  筛选前: {len(df)} 支")
        print(f"  价格: {self.config.MIN_PRICE}-{self.config.MAX_PRICE}元")
        print(f"  市值: >={self.config.MIN_MARKET_CAP/100000000:.0f}亿")
        print(f"  成交额: >={self.config.MIN_VOLUME/10000:.0f}万")
        df_before = len(df)
        df = self._filter_basic(df)
        self.stats['基础条件过滤'] = len(df)
        print(f"[OK] 通过 {len(df)} 支 (过滤 {df_before - len(df)} 支)\n")
        
        if df.empty:
            return df
        
        # 4. 技术面筛选
        print("【阶段3】技术面筛选...")
        print(f"  3月涨跌: {self.config.MIN_3M_RETURN}%~{self.config.MAX_3M_RETURN}%")
        print(f"  回撤: <={self.config.MAX_52W_DRAWDOWN}%")
        df_before = len(df)
        df = self._filter_technical(df)
        self.stats['技术面筛选'] = len(df)
        print(f"[OK] 通过 {len(df)} 支 (过滤 {df_before - len(df)} 支)\n")
        
        # 5. 评分排序
        print("【阶段4】综合评分...")
        df = self._score_and_rank(df)
        df = df.head(self.config.OUTPUT_COUNT)
        self.stats['最终入选'] = len(df)
        print(f"[OK] 最终入选: {len(df)} 支\n")
        
        return df
    
    def _calculate_indicators(self, stock: Dict, df: pd.DataFrame) -> Dict:
        """计算技术指标"""
        try:
            latest = df.iloc[-1]
            
            result = {
                '代码': stock['代码'],
                '名称': stock['名称'],
                '最新价': latest['close'],
                '换手率': latest['turn'],
                '成交额': latest['amount'],
            }
            
            # 成交额单位转换（元）
            result['成交额'] = result['成交额'] * 100 if result['成交额'] > 0 else 0
            
            # MA60
            if len(df) >= 60:
                result['MA60'] = df['close'].tail(60).mean()
            elif len(df) >= 20:
                result['MA60'] = df['close'].tail(20).mean()
            else:
                result['MA60'] = df['close'].mean()
            
            # 52周最高（或可用数据最高）
            result['52周最高'] = df['close'].max()
            
            # 3个月涨跌幅
            if len(df) >= 60:
                price_60d = df.iloc[-60]['close']
                result['3个月涨跌幅'] = ((result['最新价'] - price_60d) / price_60d) * 100 if price_60d > 0 else 0
            elif len(df) >= 20:
                price_20d = df.iloc[-20]['close']
                result['3个月涨跌幅'] = ((result['最新价'] - price_20d) / price_20d) * 100 if price_20d > 0 else 0
            else:
                result['3个月涨跌幅'] = 0
            
            # 52周回撤
            result['52周回撤'] = ((result['52周最高'] - result['最新价']) / result['52周最高']) * 100 if result['52周最高'] > 0 else 0
            
            # 平均成交额（近20日）
            if len(df) >= 20:
                result['平均成交额'] = df['amount'].tail(20).mean() * 100
            else:
                result['平均成交额'] = df['amount'].mean() * 100
            
            # 平均换手率
            if len(df) >= 20:
                result['平均换手率'] = df['turn'].tail(20).mean()
            else:
                result['平均换手率'] = df['turn'].mean()
            
            # 估算市值
            if result['换手率'] > 0:
                result['总市值'] = (result['成交额'] / result['换手率']) * 100
            else:
                result['总市值'] = 0
            
            return result
        except:
            return None
    
    def _filter_basic(self, df: pd.DataFrame) -> pd.DataFrame:
        """基础条件筛选"""
        df = df[(df['最新价'] >= self.config.MIN_PRICE) & 
                (df['最新价'] <= self.config.MAX_PRICE)]
        df = df[df['平均成交额'] >= self.config.MIN_VOLUME]
        df = df[df['总市值'] >= self.config.MIN_MARKET_CAP]
        df = df[df['平均换手率'] >= self.config.MIN_TURNOVER]
        return df.reset_index(drop=True)
    
    def _filter_technical(self, df: pd.DataFrame) -> pd.DataFrame:
        """技术面筛选"""
        df = df[(df['3个月涨跌幅'] >= self.config.MIN_3M_RETURN) & 
                (df['3个月涨跌幅'] <= self.config.MAX_3M_RETURN)]
        df = df[df['52周回撤'] <= self.config.MAX_52W_DRAWDOWN]
        df = df[df['最新价'] >= df['MA60']]
        return df.reset_index(drop=True)
    
    def _score_and_rank(self, df: pd.DataFrame) -> pd.DataFrame:
        """综合评分"""
        if df.empty:
            return df
        
        # 涨幅得分
        return_score = 1 - abs(df['3个月涨跌幅'] - 10) / 20
        df['涨幅得分'] = return_score.clip(0, 1) * 30
        
        # 回撤得分
        drawdown_score = 1 - df['52周回撤'] / self.config.MAX_52W_DRAWDOWN
        df['回撤得分'] = drawdown_score.clip(0, 1) * 30
        
        # 流动性得分
        volume_score = np.log10(df['平均成交额'] / 100_000_000).clip(0, 2) / 2
        df['成交额得分'] = volume_score * 20
        
        turnover_score = df['平均换手率'].clip(0, 10) / 10
        df['换手率得分'] = turnover_score * 20
        
        # 综合得分
        df['综合得分'] = df['涨幅得分'] + df['回撤得分'] + df['成交额得分'] + df['换手率得分']
        
        # 推荐理由
        df['推荐理由'] = df.apply(self._generate_reason, axis=1)
        
        return df.sort_values('综合得分', ascending=False)
    
    def _generate_reason(self, row) -> str:
        """生成推荐理由"""
        reasons = []
        
        ret = row['3个月涨跌幅']
        if 0 <= ret <= 10:
            reasons.append('稳健上涨')
        elif ret > 10:
            reasons.append('强势上涨')
        
        if row['52周回撤'] < 15:
            reasons.append('近高点')
        
        if row['平均换手率'] > 3:
            reasons.append('活跃')
        
        market_cap = row['总市值'] / 100_000_000
        if market_cap < 100:
            reasons.append('中小盘')
        else:
            reasons.append('中大盘')
        
        return '，'.join(reasons) if reasons else '技术面良好'
    
    def print_stats(self):
        """打印统计"""
        print("=" * 80)
        print("筛选统计".center(80))
        print("=" * 80)
        for key, value in self.stats.items():
            print(f"  {key}: {value:>6} 支")
        if self.stats['初始股票池'] > 0:
            rate = (self.stats['最终入选'] / self.stats['初始股票池']) * 100
            print(f"  筛选率: {rate:>6.2f}%")
        print("=" * 80 + "\n")


class ReportGenerator:
    """报告生成"""
    
    def __init__(self, df: pd.DataFrame, stats: dict, fetch_time: datetime):
        self.df = df
        self.stats = stats
        self.fetch_time = fetch_time
    
    def print_console(self):
        """控制台输出"""
        from tabulate import tabulate
        
        if self.df.empty:
            print("[WARN] 无筛选结果")
            return
        
        print("=" * 100)
        print(f"A股智能筛选结果 - Top {len(self.df)}".center(100))
        print("=" * 100 + "\n")
        
        cols = ['代码', '名称', '最新价', '总市值', 
                '3个月涨跌幅', '52周回撤', '平均换手率', '综合得分', '推荐理由']
        
        display_df = self.df[[c for c in cols if c in self.df.columns]].copy()
        
        # 格式化
        display_df['最新价'] = display_df['最新价'].apply(lambda x: f"{x:.2f}")
        display_df['总市值'] = display_df['总市值'].apply(lambda x: f"{x/100000000:.1f}亿")
        display_df['3个月涨跌幅'] = display_df['3个月涨跌幅'].apply(lambda x: f"{x:.1f}%")
        display_df['52周回撤'] = display_df['52周回撤'].apply(lambda x: f"{x:.1f}%")
        display_df['平均换手率'] = display_df['平均换手率'].apply(lambda x: f"{x:.1f}%")
        display_df['综合得分'] = display_df['综合得分'].apply(lambda x: f"{x:.1f}")
        
        print(tabulate(display_df, headers='keys', tablefmt='grid', showindex=False))
        print(f"\n{'=' * 100}\n")
    
    def export_excel(self, filename: str = None):
        """导出Excel"""
        if self.df.empty:
            return
        
        if not filename:
            filename = f"A股筛选_Baostock_{self.fetch_time.strftime('%Y%m%d_%H%M%S')}.xlsx"
        
        try:
            with pd.ExcelWriter(filename, engine='openpyxl') as writer:
                cols = ['代码', '名称', '最新价', '总市值', '平均成交额',
                       '3个月涨跌幅', '52周回撤', '平均换手率', 'MA60', 
                       '综合得分', '推荐理由']
                out_df = self.df[[c for c in cols if c in self.df.columns]].copy()
                
                if '总市值' in out_df.columns:
                    out_df['总市值(亿)'] = (out_df['总市值'] / 100000000).round(2)
                    out_df = out_df.drop('总市值', axis=1)
                
                out_df.to_excel(writer, sheet_name='筛选结果', index=False)
                self.df.to_excel(writer, sheet_name='详细数据', index=False)
                
                # 统计
                stats_df = pd.DataFrame(list(self.stats.items()), columns=['项目', '数量'])
                stats_df.to_excel(writer, sheet_name='统计信息', index=False)
            
            print(f"[OK] Excel已导出: {filename}\n")
        except Exception as e:
            print(f"[ERROR] 导出失败: {e}\n")


def main():
    """主函数"""
    try:
        screener = StockScreener()
        result = screener.run()
        screener.print_stats()
        
        if result.empty:
            print("[WARN] 未找到符合条件的股票，建议调整筛选参数")
            return
        
        reporter = ReportGenerator(result, screener.stats, screener.fetch_time)
        reporter.print_console()
        reporter.export_excel()
        
        print("[OK] 筛选任务完成！\n")
    except Exception as e:
        print(f"[ERROR] 程序错误: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()
